import os
from PIL import Image
import numpy as np
import threading


def select_useful_img(data_root, label_list, build_list, imp_list, grass_list):
    for i in label_list:
        label_path = data_root + '/ann_dir/train/' + i
        label = Image.open(label_path)
        w, h = label.size
        pixel_sum = w * h
        label = np.array(label)
        if (np.sum(label == 1) / pixel_sum > 0.3) :     # Search for build rate > 0.3
            if mutex.acquire():
                build_list.append(data_root + '/img_dir/train/' + i)
                mutex.release()
        elif (np.sum(label == 4) / pixel_sum > 0.3) :       # Search for imp rate > 0.3
            if mutex.acquire():
                imp_list.append(data_root + '/img_dir/train/' + i)
                mutex.release()
        elif ((np.sum(label == 3) + np.sum(label == 5)) / pixel_sum > 0.5) :         # Search for grass
            if mutex.acquire():
                grass_list.append(data_root + '/img_dir/train/' + i)
                mutex.release()



def select_count(A_build_list, A_imp_list, A_grass_list, B_build_list, B_imp_list, B_grass_list) :
    # Attention: The operation below is mainly because that Potsdam dataset is much larger than Vaihingen!
    item_min = min(len(A_build_list), min(len(A_imp_list), len(A_grass_list)))
    A_build_list = A_build_list[0:item_min]
    A_imp_list = A_imp_list[0:item_min]
    A_grass_list = A_grass_list[0:item_min]
    build_list_output = []
    imp_list_output = []
    grass_list_output = []
    for i in range(item_min) :
        build_tmp = []
        build_tmp.append(A_build_list[i])
        build_tmp.append(B_build_list[i % len(B_build_list)])
        build_list_output.append(build_tmp)

        imp_tmp = []
        imp_tmp.append(A_imp_list[i])
        imp_tmp.append(B_imp_list[i % len(B_imp_list)])
        imp_list_output.append(imp_tmp)

        grass_tmp = []
        grass_tmp.append(A_grass_list[i])
        grass_tmp.append(B_grass_list[i % len(B_grass_list)])
        grass_list_output.append(grass_tmp)

    return build_list_output, imp_list_output, grass_list_output


def get_useful_name(path, dataset_name) :
    base_name_with_no_suffix = path.split('/')[-1][0:-4]
    fn = base_name_with_no_suffix.split('_')
    if (dataset_name == 'potsdam') :
        output = dataset_name + '_' + fn[2] + '_' + fn[3] + '_' + fn[5]
    elif (dataset_name == 'vaihingen') :
        output = dataset_name + '_' + fn[3] + '_' + fn[5]
    return output


def merge_A_and_B(build_list, imp_list, grass_list) :
    for i in build_list :
        A_path = i[0]
        B_path = i[1]
        output_base_name = get_useful_name(A_path, 'potsdam') + '_' + get_useful_name(B_path, 'vaihingen') + '.png'
        img_A = Image.open(A_path)
        img_B = Image.open(B_path)
        w, h = img_A.size
        merged_image = Image.new('RGB', (2*w, h))
        merged_image.paste(img_A, (0, 0, w, h))
        merged_image.paste(img_B, (w, 0, 2 * w, h))
        merged_image.save(dstpath + '/' + output_base_name)

    for i in imp_list :
        A_path = i[0]
        B_path = i[1]
        output_base_name = get_useful_name(A_path, 'potsdam') + '_' + get_useful_name(B_path, 'vaihingen') + '.png'
        img_A = Image.open(A_path)
        img_B = Image.open(B_path)
        w, h = img_A.size
        merged_image = Image.new('RGB', (2*w, h))
        merged_image.paste(img_A, (0, 0, w, h))
        merged_image.paste(img_B, (w, 0, 2 * w, h))
        merged_image.save(dstpath + '/' + output_base_name)

    for i in grass_list :
        A_path = i[0]
        B_path = i[1]
        output_base_name = get_useful_name(A_path, 'potsdam') + '_' + get_useful_name(B_path, 'vaihingen') + '.png'
        img_A = Image.open(A_path)
        img_B = Image.open(B_path)
        w, h = img_A.size
        merged_image = Image.new('RGB', (2*w, h))
        merged_image.paste(img_A, (0, 0, w, h))
        merged_image.paste(img_B, (w, 0, 2 * w, h))
        merged_image.save(dstpath + '/' + output_base_name)

def main():
    # Do A dataset select.
    data_A_label_list = []
    for root, dirs, files in os.walk(data_root_A + '/ann_dir/train'):
        for tempName in files:
            data_A_label_list.append(tempName)
    length_A = len(data_A_label_list)
    A_build_list = []
    A_imp_list = []
    A_grass_list = []
    len_per_A = int(length_A / 8)
    t_list = []
    for i in range(7):
        t_list.append(threading.Thread(target=select_useful_img,
                                       args=(data_root_A,
                                             data_A_label_list[int(len_per_A * i): int(len_per_A * (i + 1))],
                                             A_build_list, A_imp_list, A_grass_list)))
        print('Start thread #' + str(i + 1))
    t_list.append(threading.Thread(target=select_useful_img, args=(data_root_A, data_A_label_list[int(len_per_A * 7):],
                                                                    A_build_list, A_imp_list, A_grass_list)))
    print('Start thread #' + str(8))
    for t in t_list:
        t.start()
    while True:
        if len(threading.enumerate()) == 1:
            break

    # Do B dataset select.
    data_B_label_list = []
    for root, dirs, files in os.walk(data_root_B + '/ann_dir/train'):
        for tempName in files:
            data_B_label_list.append(tempName)
    length_B = len(data_A_label_list)
    len_per_B = int(length_B / 8)
    B_build_list = []
    B_imp_list = []
    B_grass_list = []
    t_list = []
    for i in range(7):
        t_list.append(threading.Thread(target=select_useful_img,
                                       args=(data_root_B,
                                             data_B_label_list[int(len_per_B * i): int(len_per_B * (i + 1))],
                                             B_build_list, B_imp_list, B_grass_list)))
        print('Start thread #' + str(i + 1))
    t_list.append(threading.Thread(target=select_useful_img, args=(data_root_B, data_B_label_list[int(len_per_B * 7):],
                                                                   B_build_list, B_imp_list, B_grass_list)))
    print('Start thread #' + str(8))
    for t in t_list:
        t.start()
    while True:
        if len(threading.enumerate()) == 1:
            break

    build_list, imp_list, grass_list = select_count(
        A_build_list, A_imp_list, A_grass_list, B_build_list, B_imp_list, B_grass_list)
    len_per = len(build_list) / 8
    t_list = []
    for i in range(7):
        t_list.append(threading.Thread(target=merge_A_and_B,
                                       args=(build_list[int(len_per * i): int(len_per * (i + 1))],
                                             imp_list[int(len_per * i): int(len_per * (i + 1))],
                                             grass_list[int(len_per * i): int(len_per * (i + 1))])))
        print('Start thread #' + str(i + 1))
    t_list.append(threading.Thread(target=merge_A_and_B, args=(build_list[int(len_per * 7):],
                                                               imp_list[int(len_per * 7):],
                                                               grass_list[int(len_per * 7):])))
    print('Start thread #' + str(8))
    for t in t_list:
        t.start()
    while True:
        if len(threading.enumerate()) == 1:
            return



data_root_A = '/media/allen/orange/mm_dataset/mydataset_Potsdam_600_normal'
data_root_B = '/media/allen/orange/mm_dataset/vaihingen_600_normal'
dstpath = '/media/allen/orange/mm_dataset/potsdam_vaihingen_cycle_gan'


mutex = threading.Lock()

main()